package com.bpzj.util.ssh;

import com.jcraft.jsch.Channel;
import com.jcraft.jsch.JSch;
import com.jcraft.jsch.JSchException;
import com.jcraft.jsch.Session;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.io.PrintStream;
import java.nio.ByteBuffer;
import java.nio.channels.Channels;
import java.nio.channels.Pipe;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.util.ArrayList;
import java.util.List;
import java.util.logging.Level;
import java.util.logging.Logger;
import java.util.regex.Matcher;
import java.util.regex.Pattern;


/**
 * Provides similar functions as the Unix Expect tool.<br>
 * There are two ways to create an Expect object: a constructor that takes an
 * {@link InputStream} handle and {@link OutputStream} handle; or spawning a
 * process by providing a comamnd String. <br>
 * <br>
 * The API is loosely based on Perl Expect library:<br>
 * <a href="http://search.cpan.org/~rgiersig/Expect-1.15/Expect.pod">
 * http://search.cpan.org/~rgiersig/Expect-1.15/Expect.pod</a>
 *
 * <br>
 * If you are not familiar with the Tcl version of Expect, take a look at:<br>
 * <a href="http://oreilly.com/catalog/expect/chapter/ch03.html">
 * http://oreilly.com/catalog/expect/chapter/ch03.html</a> <br>
 * <br>
 * Expect uses a thread to convert InputStream to a SelectableChannel; other
 * than this, no multi-threading is used.<br>
 * A call to expect() will block for at most timeout seconds. Expect is not
 * designed to be thread-safe, in other words, do not call methods of the same
 * Expect object in different threads.
 *
 * @author Ronnie Dong
 * @version 1.1
 */
public class Expect {
  static final Logger log = Logger.getLogger(String.valueOf(Expect.class));
  /**Logging is turned off by default.*/
  static {
    log.setLevel(Level.INFO);
  }

  private OutputStream output;
  private Pipe.SourceChannel inputChannel;

  private Selector selector;

  public Expect(InputStream input, OutputStream output) {
    try {
      this.inputChannel = inputStreamToSelectableChannel(input);
      selector = Selector.open();
      inputChannel.register(selector, SelectionKey.OP_READ);
    } catch (IOException e) {
      log.throwing("Expect","Constructor", e);
      //e.printStackTrace();
    }
    this.output = output;
  }

  public static void main(String[] args) {
    try {
      JSch jsch = new JSch();
      Session session = jsch.getSession("bpzj", "10.136.214.96");
      session.setPassword("4260");
      session.setConfig("StrictHostKeyChecking", "no");
      session.connect(60 * 1000);
      Channel channel = session.openChannel("shell");
      Expect expect = new Expect(channel.getInputStream(), channel.getOutputStream());
      channel.connect();
      expect.expect("$");
      System.out.println(expect.before + expect.match);
      expect.send("ls\n");
      expect.expect("$");
      System.out.println(expect.before + expect.match);
      expect.send("exit\n");
      expect.expectEOF();
      System.out.println(expect.before);
      expect.close();
      session.disconnect();
    } catch (JSchException e) {
      e.printStackTrace();
    } catch (IOException e) {
      e.printStackTrace();
    }
  }

  /**
   * Essentially, this method converts an {@link InputStream} to a
   * {@link SelectableChannel}. A thread is created to read from the
   * InputStream, and write to a pipe. The source of the pipe is returned as
   * an input handle from which you can perform unblocking read. The thread
   * will terminate when reading EOF from InputStream, or when InputStream is
   * closed, or when the returned Channel is closed(pipe broken).
   *
   * @param input
   * @return a non-blocking Channel you can read from
   * @throws IOException
   *             most unlikely
   *
   */
  private static Pipe.SourceChannel inputStreamToSelectableChannel(
      final InputStream input) throws IOException {
    Pipe pipe = Pipe.open();
    pipe.source().configureBlocking(false);
    final OutputStream out = Channels.newOutputStream(pipe.sink());
    Thread piping = new Thread(new Runnable() {
      @Override
      public void run() {
        //LOG
        byte[] buffer = new byte[1024];
        try {
          for (int n = 0; n != -1; n = input.read(buffer)) {
            out.write(buffer, 0, n);
            if (duplicatedTo != null) {
              String toWrite = new String(buffer, 0, n);
              duplicatedTo.append(toWrite);	// no Exception will be thrown
            }
          }
          log.info("EOF from InputStream");
          input.close();		// now that input has EOF, close it.
          // other than this, do not close input
        } catch (IOException e) {
          log.throwing("Expect","IOException when piping from InputStream, "
              + "now the piping thread will end", e);
          //e.printStackTrace();
        } finally {
          try {
            log.info("closing sink of the pipe");
            out.close();
          } catch (IOException e) {
          }
        }
      }
    });
    piping.setName("Piping InputStream to SelectableChannel Thread");
    piping.setDaemon(true);
    piping.start();
    return pipe.source();
  }

  private Process process = null;
  /**
   * @return the spawned process, if this {@link Expect} object is created by
   *         spawning a process
   */
  public Process getProcess() {
    return process;
  }

  /**
   * Creates an Expect object by spawning a command.<br>
   * To Linux users, perhaps you need to use "bash -i" if you want to spawn
   * Bash.<br>
   * Note: error stream of the process is redirected to output stream.
   *
   * @param command
   * @return Expect object created using the input and output handles from the
   *         spawned process
   */
  public static Expect spawn(String command) {
    ProcessBuilder pb = new ProcessBuilder(command.split(" "));
    pb.redirectErrorStream(true);
    Process p;
    try {
      p = pb.start();
    } catch (IOException e) {
      //e.printStackTrace();
      log.throwing("Expect","Error when spawning command: " + command, e);
      return null;
    }
    Expect retv = new Expect(p.getInputStream(), p.getOutputStream());
    retv.process = p;
    return retv;
  }

  /**
   * @param str
   * Convenience method to send a string to output handle
   */
  public void send(String str) {
    this.send(str.getBytes());
  }

  /**
   * @param toWrite
   * Write a byte array to the output handle, notice flush()
   */
  public void send(byte[] toWrite) {
    //System.out.println("sending: " + bytesToPrintableString(toWrite));
    log.info("sending: " + bytesToPrintableString(toWrite));
    try {
      output.write(toWrite);
      output.flush();
    } catch (IOException e) {
      log.throwing("Expect","Error when sending bytes to output", e);
      //e.printStackTrace();
    }
  }

  private int default_timeout = 60;
  private boolean restart_timeout_upon_receive = false;
  private StringBuffer buffer = new StringBuffer();
  private boolean notransfer = false;

  /**String before the last match(if there was a match),
   *  updated after each expect() call*/
  public String before;
  /**String representing the last match(if there was a match),
   *  updated after each expect() call*/
  public String match;
  /**Whether the last match was successful,
   *  updated after each expect() call*/
  public boolean isSuccess = false;

  public static final int RETV_TIMEOUT = -1, RETV_EOF = -2,
      RETV_IOEXCEPTION = -9;

  /**
   * Convenience method, same as calling {@link #expect(int, Object...)
   * expect(default_timeout, patterns)}
   *
   * @param patterns
   * @return
   */
  public int expect(Object... patterns) {
    return expect(default_timeout, patterns);
  }

  /**
   * Convenience method, internally it constructs a List{@literal <Pattern>}
   * using the object array, and call {@link #expect(int, List) } using the
   * List. The {@link String}s in the object array will be treated as
   * literals; meanwhile {@link Pattern}s will be directly added to the List.
   * If the array contains other objects, they will be converted by
   * {@link #toString()} and then used as literal strings.
   *
   * @param patterns
   * @return
   */
  public int expect(int timeout, Object... patterns) {
    ArrayList<Pattern> list = new ArrayList<Pattern>();
    for (Object o : patterns) {
      if (o instanceof String)
        list.add(Pattern.compile(Pattern.quote((String) o))); // requires 1.5 and up
      else if (o instanceof Pattern)
        list.add((Pattern) o);
      else{
        log.info("Object " + o.toString() + " (class: "
            + o.getClass().getName() + ") is neither a String nor "
            + "a java.util.regex.Pattern, using as a literal String");
        list.add(Pattern.compile(Pattern.quote(o.toString())));
      }
    }
    return expect(timeout, list);
  }

  /**
   * Expect will wait for the input handle to produce one of the patterns in
   * the list. If a match is found, this method returns immediately;
   * otherwise, the methods waits for up to timeout seconds, then returns. If
   * timeout is less than or equal to 0 Expect will check one time to see if
   * the internal buffer contains the pattern.
   *
   * @param timeout
   *            timeout in seconds
   * @param list
   *            List of Java {@link Pattern}s used for match the internal
   *            buffer obtained by reading the InputStream
   * @return position of the matched pattern within the list (starting from
   *         0); or a negative number if there is an IOException, EOF or
   *         timeout
   */
  public int expect(int timeout, List<Pattern> list) {
    log.info("Expecting " + list);

    clearGlobalVariables();
    long endTime = System.currentTimeMillis() + (long)timeout * 1000;

    try {
      ByteBuffer bytes = ByteBuffer.allocate(1024);
      int n;
      while (true) {
        for (int i = 0; i < list.size(); i++) {
          Matcher m = list.get(i).matcher(buffer);
          if (m.find()) {
            int matchStart = m.start(), matchEnd = m.end();
            this.before = buffer.substring(0, matchStart);
            this.match = m.group();
            this.isSuccess = true;
            if(!notransfer)buffer.delete(0, matchEnd);
            return i;
          }
        }

        long waitTime = endTime - System.currentTimeMillis();
        if (restart_timeout_upon_receive)
          waitTime = timeout * 1000;
        if (waitTime <= 0) {
          log.info("Timeout when expecting " + list);
          return RETV_TIMEOUT;
        }
        //System.out.println("waiting for "+waitTime);

        selector.select(waitTime);
        //System.out.println(selector.selectedKeys().size());
        if (selector.selectedKeys().size() == 0) {
          //System.err.println("timeout!");
          //break;	//we can directly "break" here
          log.info("Timeout when expecting " + list);
          return RETV_TIMEOUT;
        }
        selector.selectedKeys().clear();
        if ((n = inputChannel.read(bytes)) == -1) {
          //System.err.println("EOF!");
          //break;
          log.info("EOF when expecting " + list);
          return RETV_EOF;
        }
        StringBuilder tmp = new StringBuilder();
        for (int i = 0; i < n; i++) {
          buffer.append((char) bytes.get(i));
          tmp.append(byteToPrintableString(bytes.get(i)));
        }
        log.info("Obtained following from InputStream: " + tmp);
        bytes.clear();

        //System.out.println(buffer);
      }
    } catch (IOException e) {
      //e.printStackTrace();
      log.throwing("Expect","IOException when selecting or reading", e);
      thrownIOE = e;
      return RETV_IOEXCEPTION;
    }

  }

  /**
   * Convenience method, internally it calls {@link #expect(int, List)
   * expect(timeout, new ArrayList&lt;Pattern&gt;())}. Given an empty list,
   * {@link #expect(int, List)} will not perform any regex matching, therefore
   * the only conditions for it to return is EOF or timeout (or IOException).
   * If EOF is detected, {@link #isSuccess} and {@link #before} are properly
   * set.
   *
   * @param timeout
   * @return same as return value of {@link #expect(int, List)}
   */
  public int expectEOF(int timeout) {
    int retv = expect(timeout, new ArrayList<Pattern>());
    if (retv == RETV_EOF) {
      this.isSuccess = true;
      this.before = this.buffer.toString();
      this.buffer.delete(0, buffer.length());
    }
    return retv;
  }
  /**Convenience method, same as calling {@link #expectEOF(int)
   * expectEOF(default_timeout)}*/
  public int expectEOF() {
    return expectEOF(default_timeout);
  }

  /**
   * Throws checked exceptions when expectEOF was not successful.
   */
  public int expectEOFOrThrow(int timeout) throws TimeoutException,
      IOException {
    int retv = expectEOF(timeout);
    if (retv == RETV_TIMEOUT)
      throw new TimeoutException();
    if (retv == RETV_IOEXCEPTION)
      throw thrownIOE;
    return retv;
  }
  /**Convenience method, same as calling {@link #expectEOF(int)
   * expectEOF(default_timeout)}*/
  public int expectEOFOrThrow() throws TimeoutException, IOException {
    return expectEOFOrThrow(default_timeout);
  }

  /**useful when calling {@link #expectOrThrow(int, Object...)}*/
  private IOException thrownIOE;

  /**
   * This method calls {@link #expect(int, Object...) expect(timeout,
   * patterns)}, and throws checked exceptions when expect was not successful.
   * Useful when you want to simplify error handling: for example, when you
   * send a series of commands to an SSH server, you expect a prompt after
   * each send, however the server may die or the prompt may take forever to
   * appear, you would want to skip the following commands if those occurred.
   * In such a case this method will be handy.
   *
   * @param timeout
   * @param patterns
   * @throws TimeoutException
   *             when expect times out
   * @throws EOFException
   *             when EOF is encountered
   * @throws IOException
   *             when there is a problem reading from the InputStream
   * @return same as {@link #expect(int, Object...) expect(timeout, patterns)}
   */
  public int expectOrThrow(int timeout, Object... patterns)
      throws TimeoutException, EOFException, IOException {
    int retv = expect(timeout, patterns);
    switch (retv) {
      case RETV_TIMEOUT:
        throw new TimeoutException();
      case RETV_EOF:
        throw new EOFException();
      case RETV_IOEXCEPTION:
        throw thrownIOE;
      default:
        return retv;
    }
  }
  /**Convenience method, same as calling {@link #expectOrThrow(int, Object...)
   * expectOrThrow(default_timeout, patterns)}*/
  public int expectOrThrow(Object... patterns) throws TimeoutException,
      EOFException, IOException {
    return expectOrThrow(default_timeout, patterns);
  }

  private void clearGlobalVariables() {
    isSuccess = false;
    match = null;
    before = null;
  }

  /**
   * The OutputStream passed to Expect constructor is closed; the InputStream
   * is not closed (there is no need to close the InputStream).<br>
   * It is suggested that this method be called after the InputStream has come
   * to EOF. For example, when you connect through SSH, send an "exit" command
   * first, and then call this method.<br>
   * <br>
   *
   * When this method is called, the thread which write to the sink of the
   * pipe will end.
   */
  public void close() {
    try {
      this.output.close();
    } catch (IOException e) {
      log.throwing("Expect","Exception when closing OutputStream", e);
      //e.printStackTrace();
    }
    try {
      this.inputChannel.close();
    } catch (IOException e) {
      log.throwing("Expect","Exception when closing input Channel", e);
      //e.printStackTrace();
    }
  }

  public int getDefault_timeout() {
    return default_timeout;
  }
  public void setDefault_timeout(int default_timeout) {
    this.default_timeout = default_timeout;
  }
  public boolean isRestart_timeout_upon_receive() {
    return restart_timeout_upon_receive;
  }
  public void setRestart_timeout_upon_receive(boolean restart_timeout_upon_receive) {
    this.restart_timeout_upon_receive = restart_timeout_upon_receive;
  }
  public void setNotransfer(boolean notransfer) {
    this.notransfer = notransfer;
  }
  public boolean isNotransfer() {
    return notransfer;
  }

  /**
   * Static method used for convert byte array to string, each byte is
   * converted to an ASCII character, if the byte represents a control
   * character, it is replaced by a printable caret notation <a
   * href="http://en.wikipedia.org/wiki/ASCII">
   * http://en.wikipedia.org/wiki/ASCII </a>, or an escape code if possible.
   *
   * @param bytes
   *            bytes to be printed
   * @return String representation of the byte array
   */
  public static String bytesToPrintableString(byte[] bytes) {
    StringBuilder sb = new StringBuilder();
    for (byte b : bytes)
      sb.append(byteToPrintableString(b));
    return sb.toString();
  }
  public static String byteToPrintableString(byte b) {
    String s = new String(new byte[] { b });
    // control characters
    if (b >= 0 && b < 32) s = "^" + (char) (b + 64);
    else if (b == 127) s = "^?";
    // some escape characters
    if (b == 9) s = "\\t";
    if (b == 10) s = "\\n";
    if (b == 13) s = "\\r";
    return s;
  }

  @SuppressWarnings("serial")
  public static class TimeoutException extends Exception{
  }
  @SuppressWarnings("serial")
  public static class EOFException extends Exception{
  }


  private static PrintStream duplicatedTo = null;
  /**
   * While performing expect operations on the InputStream provided, duplicate
   * the contents obtained from InputStream to a PrintStream (you can use
   * System.err or System.out). <b>DO NOT</b> call this function while there
   * are live Expect objects as this may cause the piping thread to end due to
   * unsynchronized code; if you need this feature, add the following to both
   * {@link #inputStreamToSelectableChannel(InputStream)} and
   * {@link #forwardInputStreamTo(PrintStream)}:
   * <pre>
   * {@code
   * 	synchronized(Expect.duplicatedTo) {...}
   * </pre>
   * @param duplicatedTo
   *            call with null if you want to turn off
   */
  public static void forwardInputStreamTo(PrintStream duplicatedTo) {
    Expect.duplicatedTo = duplicatedTo;
  }

}
